package com.wht.ai;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.wht.ai.anythingLLM.AnythingLLMDialect;
import com.wht.ai.tools.MenuFunction;
import lombok.SneakyThrows;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatResponse;
import org.noear.solon.ai.chat.dialect.ChatDialectManager;
import org.noear.solon.ai.chat.message.ChatMessage;

import java.util.Scanner;

/**
 * @author by HaiTao.Wang on 2025/2/18.
 */
public class FunctionTest {

    @SneakyThrows
    public static void main(String[] args) {

        //Ollama qwen2.5 模型
        ChatModel chatModel = ChatModel.of("http://127.0.0.1:11434/api/chat")
                .provider("ollama")
                .model("qwen2.5:14b")
                //添加工具类方法
                .globalFunctionAdd(new MenuFunction())
                .build();

        //注册anythingLLM方言
        ChatDialectManager.register(new AnythingLLMDialect());

        //构建anythingLLM ChatModel
        ChatModel anythingModel = ChatModel.of("http://localhost:3001/api/v1/workspace/solon-test/chat")
                .headerSet("Authorization", "Bearer TY5KDCS-3EQ47FB-G1NAD4N-W4YXTD6")
                .provider("anythingLLM")
                .build();

        //大模型决策关键点
        String text = "，这个问题的答案是否需要调用function tool？请回答“需要”或者“不需要”，不要回答其他额外的文字";

        System.err.println("顾客您好！ 请问您有什么需要？");
        Scanner scanner = new Scanner(System.in);
        String userInput = scanner.nextLine();
        while (StrUtil.isNotEmpty(userInput)) {

            String content = JSONUtil.createObj()
                    .set("input", JSONUtil.createObj().set("prompt", userInput).toString())
                    .toString();

            //大模型决策
            ChatResponse decision = chatModel.prompt(ChatMessage.ofUser(content + text)).call();

            //如果需要调用function tool，则调用function tool来回答
            if ("需要".equals(decision.getMessage().getContent())) {
                ChatResponse response = chatModel.prompt(userInput).call();
                System.err.println("function answer:" + response.getMessage().getContent());
                //等待用户输入
                userInput = scanner.nextLine();
                continue;
            }

            //如果不需要则用知识库来回答
            ChatResponse responseAnything = anythingModel.prompt(userInput).call();
            System.err.println("anything answer:" + responseAnything.getMessage().getContent());

            if ("q".equals(userInput)) {
                break;
            }

            //等待用户输入
            userInput = scanner.nextLine();

        }

        System.err.println("对话结束！");
    }
}
